from flask import Flask, request, jsonify
import joblib
import pandas as pd
import logging
from logging.handlers import RotatingFileHandler

app = Flask(__name__)

# 配置日志（增强：记录错误时的请求数据）
handler = RotatingFileHandler('api_server.log', maxBytes=10000, backupCount=1)
handler.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
app.logger.addHandler(handler)

try:
    model1_export = joblib.load('model1_export.pkl')
    model1_import = joblib.load('model1_import.pkl') 
    model2_export = joblib.load('model2_export.pkl')
    model2_import = joblib.load('model2_import.pkl')
    app.logger.info("模型加载成功")
except Exception as e:
    app.logger.error(f"模型加载失败: {str(e)}")
    exit(1)

def parse_time(time_str):
    try:
        year, month = map(int, time_str.split('-'))
        return {'年份': year, '月份': month}
    except Exception as e:
        app.logger.error(f"时间解析失败: {str(e)}, 输入时间: {time_str}")
        raise  # 重新抛出异常，让上层处理

@app.route('/predict/country', methods=['POST'])
def predict_country():
    data = request.get_json()
    try:
        time_info = parse_time(data['time'])
        country = data['country']
        X = pd.DataFrame([{
            '年份': time_info['年份'],
            '月份': time_info['月份'],
            '国家标签': country
        }])
        export_pred = model1_export.predict(X)[0]
        import_pred = model1_import.predict(X)[0]
        return jsonify({
            "export": round(export_pred, 2),
            "import": round(import_pred, 2)
        })
    except Exception as e:
        app.logger.error(f"国家预测失败: {str(e)}, 请求数据: {data}")  # 记录详细错误和请求数据
        return jsonify({"error": str(e)}), 400

@app.route('/predict/total', methods=['POST'])
def predict_total():
    data = request.get_json()
    try:
        time_info = parse_time(data['time'])
        X = pd.DataFrame([{
            '年份': time_info['年份'],
            '月份': time_info['月份']
        }])
        export_pred = model2_export.predict(X)[0]
        import_pred = model2_import.predict(X)[0]
        return jsonify({
            "export": round(export_pred, 2),
            "import": round(import_pred, 2)
        })
    except Exception as e:
        app.logger.error(f"总量预测失败: {str(e)}, 请求数据: {data}")  # 记录详细错误和请求数据
        return jsonify({"error": str(e)}), 400

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5001, debug=True)